package cn.lgwen.spark.ml.learning.kaggle;

import org.apache.spark.ml.classification.DecisionTreeClassificationModel;
import org.apache.spark.ml.classification.DecisionTreeClassifier;
import org.apache.spark.ml.classification.MultilayerPerceptronClassificationModel;
import org.apache.spark.ml.classification.MultilayerPerceptronClassifier;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

/**
 * 2020/3/19
 * aven.wu
 * danxieai258@163.com
 */
public class TitanicMultilayerPerceptronClassifier {
    public static void main(String[] args) {
        SparkSession spark = SparkSession
                .builder().master("local[*]")
                .appName("TitanicMultilayerPerceptronClassifier")
                .getOrCreate();

        int[] layers = new int[] {8, 4, 2};

        MultilayerPerceptronClassifier trainer = new MultilayerPerceptronClassifier()
                .setLayers(layers)
                .setLabelCol("Survived")
                //.setSeed(123L)
                .setMaxIter(10);

        Dataset<Row> trainSet = TitanicUtil.trainData(spark);
        MultilayerPerceptronClassificationModel model = trainer.fit(trainSet);
        Dataset<Row> testSet = TitanicUtil.testData(spark);
        Dataset<Row> res = model.transform(testSet);
        System.out.println(TitanicUtil.evaluate(res));

    }
}
